# %%

from conn import ac_task_col, clip_annotated_record as clip_anno_col, mongo_client
from tqdm import tqdm
import re
from collections import defaultdict
from datetime import datetime
from pymongo import InsertOne
import pytz


# %%
anno_type = 21
anno_name = "PARKING_FREE_SPACE"
query = dict(annotate_type=anno_type, status={"$in": [1, 2]})

# %%
total = ac_task_col.count_documents(query)

# %%
deleted_tasks = []
for task in tqdm(
    ac_task_col.find(query, {"task_name": 1, "clip_ids": 1, "created_at": 1}),
    total=total,
):
    if not re.search(r"[0-9]{4}-[0-9]{2}-[0-9]{2}", task["task_name"]):
        deleted_tasks.append(task)


# %%
delete_task_ids = []
delete_clip_ids = []

for task in deleted_tasks:
    delete_task_ids.append(task["_id"])
    delete_clip_ids.extend(task["clip_ids"])

print("deleted task cnt", len(delete_task_ids))
print("deleted clip cnt", len(delete_clip_ids))


# %%
def get_priority_by_date(date: str):
    if date.startswith("2025"):
        return 2
    if date.startswith("2024-11") or date.startswith("2024-12"):
        return 3
    return 5


def get_utc_time() -> datetime:
    """
    获取当前UTC时间
    """
    utc_tz = pytz.utc
    now = datetime.now()
    utc_time = now.astimezone(utc_tz)

    return utc_time


# %%
query = dict(annotate_type=anno_type, clip_id={"$in": delete_clip_ids})
# 只获取关注的部分数据
projection = {
    "clip_id": 1,
    "name": 1,
    "collect_time": 1,
    "plate_no": 1,
    "annotate_type": 1,
    "frame_cnt": 1,
    "car_type": 1,
    "supplier": 1,
}
cur_type_clip_anno_list = clip_anno_col.find(query, projection).to_list()
assert len(cur_type_clip_anno_list) == len(delete_clip_ids)
print(f"find total clip_anno: {len(cur_type_clip_anno_list)}")

exist_ac_task_cnt = ac_task_col.count_documents({"annotate_type": anno_type})
print(f"exist ac task cnt: {exist_ac_task_cnt}")

# 按照车天进行分组
vehicle_day_group = defaultdict(list)
for clip in cur_type_clip_anno_list:
    vehicle_day = str(datetime.fromtimestamp(clip["collect_time"] / 1e6).date())
    key = (clip["plate_no"], vehicle_day)
    vehicle_day_group[key].append(clip)

# 非 21 的话，10个一组，21 的话，100个一组
CHUNK = 100
ac_task_insert_ops = []
task_idx = exist_ac_task_cnt
for veh_day, clips in vehicle_day_group.items():
    print(f"vehicle day: {veh_day}, clip count: {len(clips)}")
    for ix in range(0, len(clips), CHUNK):
        task_idx += 1
        cur_task_clips = clips[ix : ix + CHUNK]
        new_task = {
            "task_name": f"[{anno_type}][{anno_name}][{veh_day[0]}][{veh_day[1]}]验收任务{task_idx}",
            "task_description": "",
            "status": 1,
            "project": 6,
            "frequency": "1Hz",
            "annotate_type": anno_type,
            "is_vis_processed": True,
            "clip_ids": [clip["clip_id"] for clip in cur_task_clips],
            "clip_ids_num": len(cur_task_clips),
            "clip_annotated_records": cur_task_clips,
            "priority": get_priority_by_date(veh_day[1]),
            "created_at": get_utc_time(),
        }
        ac_task_insert_ops.append(InsertOne(new_task))

print(len(ac_task_insert_ops))
print(ac_task_insert_ops[0])
# %%
print("ready to insert task, count: ", len(ac_task_insert_ops))
with mongo_client.start_session() as session, session.start_transaction():
    ret = ac_task_col.bulk_write(ac_task_insert_ops, ordered=False, session=session)
    assert ret.inserted_count == len(ac_task_insert_ops)
    # raise Exception("insert success")
print(f"insert task success, total: {len(ac_task_insert_ops)}")

# %%
ret = ac_task_col.delete_many({"_id": {"$in": delete_task_ids}})
print(ret.deleted_count)
